# _*_ coding : utf-8 _*_
# @Time : 2024-06-05 18:13
# @Author : haowen
# @File : forecastshars
# @Project : pyserver
import pandas as pd
import numpy as np
import akshare as ak
from statsmodels.tsa.statespace.sarimax import SARIMAX
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.holtwinters import ExponentialSmoothing
from datetime import datetime, timedelta
from pyecharts.charts import Line
from pyecharts import options as opts

def sarimax_model(close_prices):
    order = (1, 1, 1)  # (p, d, q)
    seasonal_order = (1, 1, 1, 7)  # (P, D, Q, s)
    model = SARIMAX(close_prices, order=order, seasonal_order=seasonal_order)
    results = model.fit()
    return results

def arima_model(close_prices):
    order = (5, 1, 0)  # (p, d, q)
    model = ARIMA(close_prices, order=order)
    results = model.fit()
    return results

def exponential_smoothing_model(close_prices):
    model = ExponentialSmoothing(close_prices, seasonal='add', seasonal_periods=7)
    results = model.fit()
    return results

def yuce(startData, sharesCode, model_name):
    # 读取历史股票数据
    start_date = startData
    symbol = sharesCode
    data = ak.stock_zh_a_hist(
        symbol=symbol,
        period="daily",
        start_date=start_date,
        adjust="hfq"
    )
    data.sort_values(by='日期', inplace=True)

    # 提取日期和收盘价作为特征和目标变量
    dates = pd.to_datetime(data['日期'])
    close_prices = data['收盘'].astype(float)

    # 选择模型进行拟合和预测
    if model_name == 'sarimax':# 季节性自回归积分滑动平均模型
        results = sarimax_model(close_prices)
        predictions = results.get_forecast(steps=30).predicted_mean
    elif model_name == 'arima':# 自回归积分滑动平均模型
        results = arima_model(close_prices)
        predictions = results.forecast(steps=30)
    elif model_name == 'exponential_smoothing':# 指数平滑模型
        results = exponential_smoothing_model(close_prices)
        predictions = results.forecast(steps=30)
    else:
        raise ValueError("Unsupported model name")

    # 创建一个DataFrame来保存预测结果
    future_dates = pd.date_range(start=dates.iloc[-1] + timedelta(days=1), periods=30)
    future_index = pd.DatetimeIndex(future_dates)
    predicted_df = pd.DataFrame({
        '日期': future_index,
        '预测收盘价': predictions
    })
    # 保留2位小数
    predicted_df['预测收盘价'] = predicted_df['预测收盘价'].round(2)
    # 修改字段名
    da = data[['日期', '收盘']].rename(columns={'收盘': '历史收盘价'})

    # 创建 Line 图表对象
    line_chart_lishi = Line()
    line_chart_yuce = Line()

    # 添加历史收盘价数据
    line_chart_lishi.add_xaxis(da["日期"].tolist())
    line_chart_lishi.add_yaxis("历史收盘价", da['历史收盘价'].tolist())

    # 添加预测收盘价数据 未来1个月
    line_chart_yuce.add_xaxis(predicted_df["日期"].tolist())
    line_chart_yuce.add_yaxis("预测收盘价", predicted_df['预测收盘价'].tolist(), color="orange")

    # 设置图表标题和坐标轴标签
    line_chart_lishi.set_global_opts(
        title_opts=opts.TitleOpts(title="股票收盘价历史"),
        xaxis_opts=opts.AxisOpts(name="日期"),
        yaxis_opts=opts.AxisOpts(name="历史收盘价")
    )
    line_chart_yuce.set_global_opts(
        title_opts=opts.TitleOpts(title="股票收盘价预测"),
        xaxis_opts=opts.AxisOpts(name="日期"),
        yaxis_opts=opts.AxisOpts(name="预测收盘价")
    )

    # 保存图表
    line_chart_lishi.render("../public/lishi.html")
    line_chart_yuce.render("../public/yuce.html")

# yuce("20240430", "000004", model_name='exponential_smoothing')
